/*
 * @lc app=leetcode.cn id=2558 lang=golang
 * @lcpr version=30119
 *
 * [2558] 从数量最多的堆取走礼物
 */

// @lcpr-template-start
package main

import (
	"container/heap"
	"math"
)

func main() {
	gifts := []int{25, 64, 9, 4, 100}
	k := 4
	pickGifts(gifts, k)
}

// @lcpr-template-end
// @lc code=start

type MaxHeap []int64

func (h MaxHeap) Len() int {
	return len(h)
}

func (h MaxHeap) Less(i, j int) bool {
	return h[i] > h[j]
}

func (h MaxHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *MaxHeap) Push(x interface{}) {
	*h = append(*h, x.(int64))
}

func (h *MaxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func pickGifts(gifts []int, k int) int64 {

	pq := &MaxHeap{}
	for _, v := range gifts {
		*pq = append(*pq, int64(v))
	}
	heap.Init(pq)

	for i := 0; i < k; i++ {
		v := heap.Pop(pq).(int64)
		v = int64(math.Sqrt(float64(v)))
		heap.Push(pq, v)
	}

	var res int64
	for len(*pq) > 0 {
		res += int64(heap.Pop(pq).(int64))
	}
	return res
}

// @lc code=end

/*
// @lcpr case=start
// [25,64,9,4,100]\n4\n
// @lcpr case=end

// @lcpr case=start
// [1,1,1,1]\n4\n
// @lcpr case=end

*/
